import numpy as np
import scipy.special as sp
import matplotlib.pyplot as plt
from scipy.integrate import quad

# ====================== Parameters ======================
q = 1.0
a = 0.15
R = 25.0
num_k = 6000
num_r = 6000

# ====================== Bessel Function Setup ======================
alpha_n = sp.jn_zeros(0, num_k)
k = alpha_n / R
J1_alpha = sp.j1(alpha_n)

# ====================== Analytical Solution ======================
def analytical_coefficients(k, q, a, R, alpha_n, J1_alpha):
    F = np.zeros_like(k)
    for m in range(len(k)):
        alpha_m = alpha_n[m]
        scaled_alpha = alpha_m * a / R
        numerator = 2 * a * sp.j1(scaled_alpha)
        denominator = alpha_m * R * J1_alpha[m]**2
        F[m] = q * numerator / denominator
    return F

F_analytic = analytical_coefficients(k, q, a, R, alpha_n, J1_alpha)

# ====================== Numerical Transform ======================
# Define load function (scalar-friendly)
def f_load(r):
    return q if r <= a else 0.0

# Vectorized version for plotting
r_plot = np.linspace(0, R, num_r)
f_analytical = np.vectorize(f_load)(r_plot)

# Adaptive numerical integration
F_numerical = np.zeros_like(k)
for m in range(len(k)):
    integrand = lambda r: f_load(r) * sp.j0(k[m]*r) * r
    F_numerical[m], _ = quad(integrand, 0, R, points=[a], limit=200)

# Apply normalization
F_numerical_normalized = F_numerical * (2 / (R**2 * J1_alpha**2))

# ====================== Inverse Transform ======================
f_reconstructed = np.zeros_like(r_plot)
for i, r_val in enumerate(r_plot):
    kr = k * r_val
    f_reconstructed[i] = np.sum(F_analytic * sp.j0(kr))

# ====================== Visualization ======================
plt.figure(figsize=(9, 3), dpi=300)

plt.subplot(1, 2, 1)
plt.plot(r_plot, f_analytical, 'b-', label='Original')
plt.plot(r_plot, f_reconstructed, 'r--', label='Reconstructed')
plt.xlim(0, 2*a)
plt.xlabel('Radial Position (m)', fontsize=10)
plt.ylabel('$p_r$ (Pa)', fontsize=10)
plt.title('Spatial Domain Comparison with Fourier-Bessel series', fontsize=10)
plt.legend(loc='lower left')
plt.grid(True)

plt.subplot(1, 2, 2)
plt.semilogy(k[:2000], np.abs(F_analytic[:2000]), 'b-', label='Analytical')
plt.semilogy(k[:2000], np.abs(F_numerical_normalized[:2000]), 'r--', label='Numerical')
plt.xlabel('Wavenumber (1/m)', fontsize=10)
plt.ylabel(r"|$\hat{p}_r$|", fontsize=10)
plt.title('Wavenumber Domain Comparison', fontsize=10)
plt.legend(loc='lower left')
plt.grid(True)

plt.tight_layout()
plt.savefig('Fourier_Bessel_Series.png', bbox_inches='tight')
plt.show()

# ====================== Validation ======================
print("Boundary Verification:")
print(f"J0(k[0]*R) = {sp.j0(k[0]*R):.2e}")
print(f"J0(k[-1]*R) = {sp.j0(k[-1]*R):.2e}")

# Energy calculation (fixed)
energy_r = quad(lambda r: f_load(r)**2 * r, 0, R)[0]
energy_k = np.sum(F_analytic**2 * (R**2/2) * J1_alpha**2)
print(f"\nEnergy Ratio: {energy_k/energy_r:.8f}")